#include "stdafx.h"
#include "HookFunctions.h"
#include "HookDll.h"
#include "..\Common\Common.h"
#include "MemOrHandleHook.h"
#include "DataDef.h"


extern bool g_bHooked;
extern CArrayEx<HOOKFUNCDESC,HOOKFUNCDESC&> g_arrHookedFunctions;

//Memory
HeapAllocDef   pOrgHeapAlloc = 0;
HeapFreeDef	  pOrgHeapFree= 0;
static HeapReAllocDef pOrgHeapReAlloc= 0;

static VirtualAllocDef pOrgVirtualAlloc = 0;
static VirtualFreeDef  pOrgVirtualFree = 0;
static VirtualAllocExDef pOrgVirtualAllocEx = 0;
static VirtualFreeExDef  pOrgVirtualFreeEx = 0;
static CoTaskMemAllocDef	pOrgCoTaskMemAlloc = 0;
static CoTaskMemReallocDef	pOrgCoTaskMemRealloc = 0;
static CoTaskMemFreeDef		pOrgCoTaskMemFree = 0;
static MapViewOfFileDef		pOrgMapViewOfFile = 0;
static MapViewOfFileExDef	pOrgMapViewOfFileEx = 0;
static UnmapViewOfFileDef	pOrgUnmapViewOfFile = 0;
static NtMapViewOfSectionDef	pOrgNtMapViewOfSection = 0;
static NtUnmapViewOfSectionDef	pOrgNtUnmapViewOfSection = 0;

//Memory
	   HeapCreateDef pOrgHeapCreate = 0;
	   HeapDestroyDef pOrgHeapDestroy = 0;
	   GlobalFreeDef pOrgGlobalFree = 0;
	   LocalFreeDef pOrgLocalFree = 0;

extern LocalReAllocDef pOrgLocalReAlloc ;
extern GlobalAllocDef pOrgGlobalAlloc ;
extern LocalAllocDef pOrgLocalAlloc ;
extern GlobalReAllocDef pOrgGlobalReAlloc ;


LPVOID WINAPI MyHeapAlloc( IN HANDLE hHeap,
                           IN DWORD dwFlags,
                           IN SIZE_T dwBytes )
{
    LPVOID lMem =  pOrgHeapAlloc( hHeap, dwFlags, dwBytes );
    CreateCallStack( lMem, dwBytes );
    return lMem;
}

LPVOID WINAPI MyHeapReAlloc( HANDLE hHeap,
                           DWORD dwFlags,
                           LPVOID lpMem,
                           SIZE_T dwBytes )
{
    LPVOID lpNewMem = pOrgHeapReAlloc( hHeap, dwFlags, lpMem, dwBytes );
    try
    {
  
        if( g_bHooked  )
        {
			RemoveCallStack(lpMem);

			CreateCallStack( lpNewMem, dwBytes );
        }
    }
    catch (...)
    {
    }
    return lpNewMem;
}

LPVOID WINAPI MyVirtualAlloc( LPVOID lpAddress,
                              SIZE_T dwSize, DWORD flAllocationType,
                              DWORD flProtect )
{
    LPVOID lMem =  pOrgVirtualAlloc( lpAddress, dwSize, flAllocationType, flProtect );
    CreateCallStack( lMem, dwSize );
    return lMem;
}

BOOL WINAPI MyVirtualFree( LPVOID lpAddress, SIZE_T dwSize, DWORD dwFreeType )
{
    RemoveCallStack( lpAddress );
    return pOrgVirtualFree( lpAddress, dwSize, dwFreeType );
}

LPVOID WINAPI MyVirtualAllocEx( HANDLE hProcess, LPVOID lpAddress,
                                SIZE_T dwSize, DWORD flAllocationType,
                                DWORD flProtect )
{
    LPVOID lMem =  pOrgVirtualAllocEx( hProcess, lpAddress, dwSize, flAllocationType, flProtect );
    CreateCallStack( lMem, dwSize );
    return lMem;
}

BOOL WINAPI MyVirtualFreeEx( HANDLE hProcess, LPVOID lpAddress, SIZE_T dwSize, DWORD dwFreeType )
{
    RemoveCallStack( lpAddress );
    return pOrgVirtualFreeEx( hProcess, lpAddress, dwSize, dwFreeType );
}

BOOL WINAPI MyHeapFree(  HANDLE hHeap,  DWORD dwFlags,  LPVOID lpMem )
{
    RemoveCallStack( lpMem );
    return pOrgHeapFree( hHeap, dwFlags, lpMem );
}

HANDLE WINAPI MyGlobalFree( HGLOBAL hMem )
{
    HANDLE hHandle = pOrgGlobalFree( hMem );
    if ( hHandle == NULL )
        RemoveCallStack( hMem );
    return hHandle;

}

HLOCAL WINAPI MyLocalFree(HLOCAL hMem )
{
    HLOCAL hHandle = pOrgLocalFree(hMem );
    if ( hHandle == NULL )
        RemoveCallStack( hMem );
    return hHandle;
}

LPVOID WINAPI MyCoTaskMemAlloc( SIZE_T cb)
{
    LPVOID lpMem = pOrgCoTaskMemAlloc( cb );
    CreateCallStack( lpMem, cb );
    return lpMem;
}

LPVOID WINAPI MyCoTaskMemRealloc(LPVOID pv, SIZE_T cb)
{
    LPVOID lpMem = pOrgCoTaskMemRealloc(pv, cb );
    if( lpMem )
    {
        CreateCallStack( lpMem, cb );
        RemoveCallStack( pv );
    }
    return lpMem;
}

void   WINAPI MyCoTaskMemFree( LPVOID pv )
{
    RemoveCallStack( pv );
}


LPVOID WINAPI MyMapViewOfFileEx( HANDLE hFileMappingObject, DWORD dwDesiredAccess,
    DWORD dwFileOffsetHigh, DWORD dwFileOffsetLow, SIZE_T dwNumberOfBytesToMap, LPVOID lpBaseAddress )
{
    LPVOID lMem = pOrgMapViewOfFileEx( hFileMappingObject, dwDesiredAccess, dwFileOffsetHigh,
    dwFileOffsetLow, dwNumberOfBytesToMap, lpBaseAddress);
    CreateCallStack( lMem, dwNumberOfBytesToMap );
    return lMem;
}

LPVOID WINAPI MyMapViewOfFile( HANDLE hFileMappingObject, DWORD dwDesiredAccess,
    DWORD dwFileOffsetHigh, DWORD dwFileOffsetLow, SIZE_T dwNumberOfBytesToMap )
{
    LPVOID lMem = pOrgMapViewOfFile( hFileMappingObject, dwDesiredAccess, dwFileOffsetHigh,
    dwFileOffsetLow, dwNumberOfBytesToMap);
    CreateCallStack( lMem, dwNumberOfBytesToMap );
    return lMem;
}

INT MyNtMapViewOfSection( HANDLE SectionHandle, HANDLE ProcessHandle, PVOID *BaseAddress,
    ULONG_PTR ZeroBits, SIZE_T CommitSize, PLARGE_INTEGER SectionOffset, PSIZE_T ViewSize,
    SECTION_INHERIT InheritDisposition, ULONG AllocationType, ULONG Win32Protect )
{
    INT s = pOrgNtMapViewOfSection( SectionHandle, ProcessHandle, BaseAddress, ZeroBits, CommitSize,
    SectionOffset, ViewSize, InheritDisposition, AllocationType, Win32Protect );
    CreateCallStack( BaseAddress, *ViewSize );
    return s;
}

INT MyNtUnmapViewOfSection( HANDLE ProcessHandle, PVOID BaseAddress )
{
    INT s = pOrgNtUnmapViewOfSection( ProcessHandle, BaseAddress );
    RemoveCallStack( (LPVOID)BaseAddress );
    return s;
}

BOOL WINAPI MyUnmapViewOfFile( LPCVOID lpBaseAddress )
{
    RemoveCallStack( (LPVOID)lpBaseAddress );
    return pOrgUnmapViewOfFile( lpBaseAddress );
}

//Memory
HANDLE WINAPI MyHeapCreate( DWORD flOptions, SIZE_T dwInitialSize, SIZE_T dwMaximumSize )
{
    HANDLE hHandle = pOrgHeapCreate( flOptions,  dwInitialSize,  dwMaximumSize );
    CreateCallStack( hHandle, TYPE_MEMORY_HANDLE );
    return hHandle;
}
BOOL   WINAPI MyHeapDestroy(HANDLE hHeap )
{
    BOOL   bRet = pOrgHeapDestroy(hHeap );
    if( bRet )
        RemoveCallStack( hHeap );
    return bRet;

}

const int MEMORY_FUN_NO = 19;
void HookMemAlloc()
{
    g_arrHookedFunctions.SetSize( MEMORY_FUN_NO ); 
    int nIndex = 0;
    HMODULE Kernel32 = GetModuleHandle( "Kernel32.dll" );
    HMODULE Ole32 = GetModuleHandle( "ole32.dll" );
    HMODULE ntdll = GetModuleHandle( "ntdll.dll" );

    REGISTER_HOOK(HeapAlloc, Kernel32 );
    REGISTER_HOOK(HeapFree,Kernel32 );
    REGISTER_HOOK(HeapReAlloc,Kernel32 );
    REGISTER_HOOK(VirtualAlloc,Kernel32 );    
    REGISTER_HOOK(VirtualFree,Kernel32 );
    REGISTER_HOOK(VirtualAllocEx,Kernel32 );    
    REGISTER_HOOK(VirtualFreeEx,Kernel32 );
    REGISTER_HOOK(GlobalAlloc,Kernel32 );
    REGISTER_HOOK(GlobalReAlloc,Kernel32 );
    REGISTER_HOOK(GlobalFree,Kernel32 );
    REGISTER_HOOK(LocalAlloc,Kernel32 );
    REGISTER_HOOK(LocalReAlloc,Kernel32 );
    REGISTER_HOOK(LocalFree,Kernel32 );
    REGISTER_HOOK(MapViewOfFile,Kernel32 );
    REGISTER_HOOK(MapViewOfFileEx,Kernel32 );
    REGISTER_HOOK(UnmapViewOfFile,Kernel32 );
    REGISTER_HOOK(CoTaskMemAlloc,Ole32 );
    REGISTER_HOOK(CoTaskMemRealloc,Ole32 );
    REGISTER_HOOK(CoTaskMemFree,Ole32 );
    ///*REGISTER_HOOK( NtMapViewOfSection, ntdll);
    //REGISTER_HOOK( NtUnmapViewOfSection, ntdll );*/

    ASSERT( nIndex == MEMORY_FUN_NO );
    HookDynamicLoadedFun( nIndex, g_arrHookedFunctions.GetData());
}